use crate::{
    ast::{Expression, Type, UnaryOp},
    error::{CompilerError, Result},
    llvm_backend::LLVMBackend,
};
use inkwell::{types::BasicType, values::BasicValueEnum, AddressSpace};

/// Pointer origin tracking for memory management
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum PointerOrigin {
    /// Pointer from heap allocation (&Type{})
    Heap,
    /// Pointer from stack reference (&variable)
    Stack,
    /// Unknown origin (parameter, return value, complex expression)
    Unknown,
}

impl<'ctx> LLVMBackend<'ctx> {
    /// Cleanup stack slices and reference counted variables before function exit
    pub(super) fn cleanup_stack_slices(&mut self) -> Result<()> {
        // 引用计数现在通过内联LLVM IR实现，不需要外部函数

        // For each tracked variable (slices and reference counted)
        for var_name in &self.stack_slices.clone() {
            // Get the variable
            let var = self.variable_map.get(var_name).ok_or_else(|| {
                CompilerError::internal(&format!("Variable not found: {}", var_name))
            })?;

            if let BasicValueEnum::PointerValue(var_alloca) = var {
                // Get the variable type
                let var_type = self.variable_type_map.get(var_name).ok_or_else(|| {
                    CompilerError::internal(&format!("Type not found for variable: {}", var_name))
                })?;

                match var_type {
                    Type::Vec(vector_type) => {
                        // Get free function for this specific case
                        let free_func = self
                            .function_map
                            .get("free")
                            .ok_or_else(|| CompilerError::internal("free function not found"))?;

                        // Handle slice cleanup
                        let element_llvm_type =
                            self.nrc_type_to_llvm_type(&Some(*vector_type.element_type.clone()))?;

                        // Build slice struct type
                        let slice_struct_type = self.context.struct_type(
                            &[
                                element_llvm_type.ptr_type(AddressSpace::default()).into(),
                                self.context.i64_type().into(),
                                self.context.i64_type().into(),
                            ],
                            false,
                        );

                        // Load the slice from the alloca
                        let slice_value = self
                            .builder
                            .build_load(
                                slice_struct_type,
                                *var_alloca,
                                &format!("load_{}_for_cleanup", var_name),
                            )?
                            .into_struct_value();

                        // Extract the ptr field (field 0)
                        let array_ptr = self
                            .builder
                            .build_extract_value(
                                slice_value,
                                0,
                                &format!("{}_array_ptr", var_name),
                            )?
                            .into_pointer_value();

                        // Check if array pointer is null before freeing
                        let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
                        let array_ptr_i8 = self
                            .builder
                            .build_bitcast(
                                array_ptr,
                                i8_ptr_type,
                                &format!("{}_array_ptr_i8", var_name),
                            )?
                            .into_pointer_value();

                        // Create basic blocks for null check
                        let parent_fn = self
                            .builder
                            .get_insert_block()
                            .and_then(|b| b.get_parent())
                            .ok_or_else(|| {
                                CompilerError::internal("No parent function for cleanup")
                            })?;

                        let free_block = self
                            .context
                            .append_basic_block(parent_fn, &format!("free_{}_array", var_name));
                        let continue_block = self.context.append_basic_block(
                            parent_fn,
                            &format!("continue_{}_cleanup", var_name),
                        );

                        // Build null check: compare pointer with null
                        let ptr_as_int = self.builder.build_ptr_to_int(
                            array_ptr_i8,
                            self.context.i64_type(),
                            "ptr_as_int",
                        )?;
                        let null_as_int = self.context.i64_type().const_int(0, false);
                        let is_null = self.builder.build_int_compare(
                            inkwell::IntPredicate::EQ,
                            ptr_as_int,
                            null_as_int,
                            "is_null_ptr",
                        )?;

                        // Branch based on null check: if null, skip to continue; otherwise free
                        let _ = self.builder.build_conditional_branch(
                            is_null,
                            continue_block,
                            free_block,
                        );

                        // Free block: free the array (only if not null)
                        self.builder.position_at_end(free_block);

                        // Call free on the internal array
                        let _ = self.builder.build_call(
                            *free_func,
                            &[array_ptr_i8.into()],
                            &format!("free_{}_array", var_name),
                        );
                        let _ = self.builder.build_unconditional_branch(continue_block);

                        // Continue block: cleanup done, continue with next variable
                        self.builder.position_at_end(continue_block);
                    }
                    Type::Rc(rc_type) => {
                        let rc_type = rc_type.clone();
                        let llvm_type =
                            self.nrc_type_to_llvm_type(&Some(Type::Rc(rc_type.clone())))?;
                        let stored_value = self.builder.build_load(
                            llvm_type,
                            *var_alloca,
                            &format!("load_{}_ptr", var_name),
                        )?;
                        let rc_ptr = self.rc_block_pointer_from_value(stored_value, &rc_type)?;
                        self.generate_rc_decrement(rc_ptr.into())?;
                    }
                    _ => {
                        // Skip other types
                        continue;
                    }
                }
            }
        }

        Ok(())
    }

    /// Generate memory allocation code
    #[allow(dead_code)]
    pub(super) fn generate_malloc(
        &mut self,
        size: &BasicValueEnum<'ctx>,
    ) -> Result<BasicValueEnum<'ctx>> {
        let malloc_func = self
            .function_map
            .get("__malloc")
            .ok_or_else(|| CompilerError::internal("__malloc function not found"))?;

        // 将大小转换为i64类型
        let size_i64 = match size {
            BasicValueEnum::IntValue(int_val) => {
                if int_val.get_type().get_bit_width() == 64 {
                    *int_val
                } else {
                    self.builder
                        .build_int_cast(*int_val, self.context.i64_type(), "size_i64")?
                }
            }
            _ => return Err(CompilerError::internal("Invalid size type for malloc")),
        };

        let ptr = self
            .builder
            .build_call(*malloc_func, &[size_i64.into()], "malloc_call")?;
        Ok(ptr.try_as_basic_value().left().unwrap().into())
    }

    /// Generate memory deallocation code
    #[allow(dead_code)]
    pub(super) fn generate_free(
        &mut self,
        ptr: &BasicValueEnum<'ctx>,
        size: &BasicValueEnum<'ctx>,
    ) -> Result<()> {
        if let BasicValueEnum::PointerValue(ptr_value) = ptr {
            let free_func = self
                .function_map
                .get("__free")
                .ok_or_else(|| CompilerError::internal("__free function not found"))?;

            // 将指针转换为i8*类型（通用指针类型）
            let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
            let generic_ptr = self
                .builder
                .build_bitcast(*ptr_value, i8_ptr_type, "generic_ptr")?;

            // 将大小转换为i64类型
            let size_i64 = match size {
                BasicValueEnum::IntValue(int_val) => {
                    if int_val.get_type().get_bit_width() == 64 {
                        *int_val
                    } else {
                        self.builder.build_int_cast(
                            *int_val,
                            self.context.i64_type(),
                            "size_i64",
                        )?
                    }
                }
                _ => return Err(CompilerError::internal("Invalid size type for free")),
            };

            let _ = self.builder.build_call(
                *free_func,
                &[generic_ptr.into(), size_i64.into()],
                "free_call",
            );
        }
        Ok(())
    }

    /// Cleanup all manually allocated variables
    pub(super) fn cleanup_manual_allocated_variables(&mut self) -> Result<()> {
        // 手动内存管理：在作用域结束时，程序员需要手动释放内存
        // 这里我们只是清空变量映射，不进行自动释放
        self.variable_map.clear();
        Ok(())
    }

    /// Detect the origin of a pointer from an expression
    #[allow(dead_code)]
    pub(super) fn detect_pointer_origin(&self, expr: &Expression) -> Option<PointerOrigin> {
        match expr {
            // &Type{} - heap allocation
            Expression::Unary(unary) if unary.operator == UnaryOp::AddrOf => {
                match &*unary.operand {
                    // &Type{...} - heap allocation with struct literal
                    Expression::Struct(_) => Some(PointerOrigin::Heap),
                    // &variable - stack reference
                    Expression::Variable(_) => Some(PointerOrigin::Stack),
                    // Other cases like &(complex_expr) - unknown
                    _ => Some(PointerOrigin::Unknown),
                }
            }
            // Direct variable reference - inherit origin from that variable
            Expression::Variable(var_name) => self.pointer_origin_map.get(var_name).copied(),
            // Other expressions don't create pointers or have unknown origin
            _ => Some(PointerOrigin::Unknown),
        }
    }

    /// Generate deep free for slice (frees internal array, then slice structure)
    /// This is called when freeing a pointer to a slice: free(slice_ptr)
    pub(super) fn generate_deep_free_slice(
        &mut self,
        var_name: &str,
    ) -> Result<BasicValueEnum<'ctx>> {
        // Get the slice pointer variable
        let slice_ptr_var = self
            .variable_map
            .get(var_name)
            .ok_or_else(|| CompilerError::internal(&format!("Variable not found: {}", var_name)))?
            .clone();

        // Load the actual slice pointer (it's stored in an alloca)
        let slice_ptr = if let BasicValueEnum::PointerValue(ptr_alloca) = slice_ptr_var {
            // Get the type being pointed to
            let var_type = self.variable_type_map.get(var_name).ok_or_else(|| {
                CompilerError::internal(&format!("Type not found for variable: {}", var_name))
            })?;

            // Determine the LLVM type for the slice
            let slice_llvm_type = if let Type::Reference(ref_type) = var_type {
                if let Type::Vec(vector_type) = &*ref_type.referenced_type {
                    let element_llvm_type =
                        self.nrc_type_to_llvm_type(&Some(*vector_type.element_type.clone()))?;
                    self.context.struct_type(
                        &[
                            element_llvm_type.ptr_type(AddressSpace::default()).into(),
                            self.context.i64_type().into(),
                            self.context.i64_type().into(),
                        ],
                        false,
                    )
                } else {
                    return Err(CompilerError::internal("Expected pointer to slice type"));
                }
            } else {
                return Err(CompilerError::internal("Expected pointer type"));
            };

            // Load the slice pointer from the alloca
            self.builder
                .build_load(
                    slice_llvm_type.ptr_type(AddressSpace::default()),
                    ptr_alloca,
                    "load_slice_ptr",
                )?
                .into_pointer_value()
        } else {
            return Err(CompilerError::internal(
                "Expected pointer value for slice variable",
            ));
        };

        // Load the slice structure from the pointer
        let var_type = self.variable_type_map.get(var_name).ok_or_else(|| {
            CompilerError::internal(&format!("Type not found for variable: {}", var_name))
        })?;

        let element_llvm_type = if let Type::Reference(ref_type) = var_type {
            if let Type::Vec(vector_type) = &*ref_type.referenced_type {
                self.nrc_type_to_llvm_type(&Some(*vector_type.element_type.clone()))?
            } else {
                return Err(CompilerError::internal("Expected pointer to slice type"));
            }
        } else {
            return Err(CompilerError::internal("Expected pointer type"));
        };

        let slice_struct_type = self.context.struct_type(
            &[
                element_llvm_type.ptr_type(AddressSpace::default()).into(),
                self.context.i64_type().into(),
                self.context.i64_type().into(),
            ],
            false,
        );

        let slice_struct = self
            .builder
            .build_load(slice_struct_type, slice_ptr, "load_slice_struct")?
            .into_struct_value();

        // Extract the ptr field (field 0) from the slice structure
        let array_ptr = self
            .builder
            .build_extract_value(slice_struct, 0, "slice_array_ptr")?
            .into_pointer_value();

        // Free the internal array first (if not null)
        let free_func = self
            .function_map
            .get("free")
            .ok_or_else(|| CompilerError::internal("free function not found"))?;

        let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        let array_ptr_i8 = self
            .builder
            .build_bitcast(array_ptr, i8_ptr_type, "array_ptr_i8")?
            .into_pointer_value();

        // Call free on the internal array
        let _ = self
            .builder
            .build_call(*free_func, &[array_ptr_i8.into()], "free_slice_array");

        // Now free the slice structure itself
        let slice_ptr_i8 = self
            .builder
            .build_bitcast(slice_ptr, i8_ptr_type, "slice_ptr_i8")?
            .into_pointer_value();

        let _ = self
            .builder
            .build_call(*free_func, &[slice_ptr_i8.into()], "free_slice_struct");

        // Return a dummy value (free returns void, but we need to return something)
        Ok(self.context.i32_type().const_int(0, false).into())
    }
}
